# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import logging

from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import (
    _RemoteDatasetLoader,
)
from pyrit.models import SeedDataset, SeedPrompt

logger = logging.getLogger(__name__)


class _MultilingualVulnerabilityDataset(_RemoteDatasetLoader):
    """
    Loader for the Multilingual Vulnerability dataset.

    This dataset contains multilingual prompts demonstrating LLM vulnerabilities,
    labeled by type. From "A Framework to Assess Multilingual Vulnerabilities of LLMs".

    Reference: https://arxiv.org/pdf/2503.13081
    GitHub: https://github.com/CarsonDon/Multilingual-Vuln-LLMs
    """

    def __init__(
        self,
        *,
        source: str = "https://raw.githubusercontent.com/CarsonDon/Multilingual-Vuln-LLMs/main/prompts/allprompt.csv",
    ):
        """
        Initialize the Multilingual Vulnerability dataset loader.

        Args:
            source: URL to the CSV file. Defaults to the official repository.
        """
        self.source = source

    @property
    def dataset_name(self) -> str:
        """Return the dataset name."""
        return "multilingual_vulnerability"

    async def fetch_dataset(self, *, cache: bool = True) -> SeedDataset:
        """
        Fetch Multilingual Vulnerability dataset and return as SeedDataset.

        Args:
            cache: Whether to cache the fetched dataset. Defaults to True.

        Returns:
            SeedDataset: A SeedDataset containing the multilingual vulnerability prompts.
        """
        logger.info(f"Loading Multilingual Vulnerability dataset from {self.source}")

        # Use standard caching mechanism
        examples = self._fetch_from_url(
            source=self.source,
            source_type="public_url",
            cache=cache,
        )

        seed_prompts = [
            SeedPrompt(
                value=item["en"],
                data_type="text",
                dataset_name=self.dataset_name,
                harm_categories=[item["type"]],
                description=(
                    "Dataset from 'A Framework to Assess Multilingual Vulnerabilities of LLMs'. "
                    "Multilingual prompts demonstrating LLM vulnerabilities, labeled by type. "
                    "Paper: https://arxiv.org/pdf/2503.13081"
                ),
                authors=[
                    "Likai Tang",
                    "Niruth Bogahawatta",
                    "Yasod Ginige",
                    "Jiarui Xu",
                    "Shixuan Sun",
                    "Surangika Ranathunga",
                    "Suranga Seneviratne",
                ],
                source="https://github.com/CarsonDon/Multilingual-Vuln-LLMs",
            )
            for item in examples
        ]

        logger.info(f"Successfully loaded {len(seed_prompts)} prompts from Multilingual Vulnerability dataset")

        return SeedDataset(seeds=seed_prompts, dataset_name=self.dataset_name)
